from __future__ import annotations

import datetime
import json
from dataclasses import dataclass
from collections import OrderedDict

from typing import List, Optional, Type, Tuple, Union
from numpy.typing import DTypeLike

import peewee
from peewee import DateTimeField
from peewee import fn

from common.typing_utils import numeric
from common.geometry.geom import Geom
from common.geometry.srs import mgrs_tile_to_epsg
from common.datetime_utils import datetimestr_to_datetime
from common.geoimage.virtual_raster import VRTScene

from catalog.utils import interpret_protocol
from catalog.models.base_model import BaseItem


@dataclass
class BandDefinition:
    """Band Definition

    Attributes
    -----------
    band_name : str
        Name of the band
    asset_name : str
        Name of the asset where the band is in
    in_asset_index : int
        Index of the band within the asset
    source_resolution : numeric
        Resolution of the source data for the band
    nodata_val : numeric
        Nodata value of the band
    scale : numeric
        Scale of the band
    offset : numeric
        Offset of the band
    dtype : DTypeLike
        Data type of the band
    linear_scale : Optional[Tuple[numeric, numeric]]
        Linear scale of the band used in display,
        values in the linear_scale's [min, max] is to be scaled to [0, 255].
    """

    band_name: str
    asset_name: str
    in_asset_index: int
    source_resolution: numeric
    nodata_val: numeric
    scale: numeric
    offset: numeric
    dtype: DTypeLike
    linear_scale: Optional[Tuple[numeric, numeric]]


class BaseSTACItem(BaseItem):
    """Base class for all STAC Items managed by the catalog.

    Attributes
    ----------
    start_datetime : DateTimeField
        Start datetime of the item, indexed.
    end_datetime : DateTimeField
        End datetime of the item, indexed.
        The combination of `start_datetime` and `end_datetime` determines the acquisition time of the item.
    processed : DateTimeField
        When data processed by original provider
    published : DateTimeField
        When data published by original provider
    indexed : DateTimeField
        When data were seen by PAII
    transferred : DateTimeField
        When data were transferred to PAII
    archived : DateTimeField
        When data were processed and archived by PAII
    created : DateTimeField
        When this data entry as a STAC json created by PAII
    updated : DateTimeField
        When data entry last updated by PAII
    """

    # timestamps for life cycle management
    start_datetime = DateTimeField(index=True)
    end_datetime = DateTimeField(index=True)
    processed = DateTimeField()
    published = DateTimeField()
    indexed = DateTimeField()
    transferred = DateTimeField()
    archived = DateTimeField()
    created = DateTimeField()
    updated = DateTimeField()

    @classmethod
    def construct_query(
        cls: Type[BaseSTACItem],
        spatial: Optional[Geom] = None,
        temporal: Optional[
            Tuple[Optional[datetime.datetime], Optional[datetime.datetime]]
        ] = None,
        include_invalidated: bool = False,
    ) -> Tuple[peewee.Query, int]:
        """Construct a query for the given spatial and temporal constraints.

        Parameters
        ----------
        spatial : Optional[Geom], optional
            Area of interest represented by a Geom object, by default None not to filter by spatial area.
        temporal : Optional[Tuple[Optional[datetime.datetime], Optional[datetime.datetime]] ], optional
            Temporal range of interest, as a tuple of start and end datetime, or datetime string,
            by default None not to filter by temporal range.
        include_invalidated : bool, optional
            Whether to include invalidated items in the query result, by default False not to include.

        Returns
        -------
        Tuple[peewee.Query, int]
            A query object that can be further tuned or executed, and the number of criterion that were applied.
        """
        # NOTE This class method provides the most common and basic query
        # including filter out the invalidated items via `include_invalidated`
        # and the "spatial, temporal" query terms
        # and return a Tuple[peewee.Query, int]
        # with the first item a query object
        # and the last item an integer, as the number of query terms applied
        # It is expected to be override by each actual model sub-class, to perform its own query
        # typically accept query in the format of "spatial, temporal" and properties queries as kwargs
        # then the sub-class method could call this to construct the base query,
        # and further apply more query terms, at the end return the query, n_conditions

        query = cls.select()
        if not include_invalidated:
            query = query.where(cls.valid_until.is_null())

        n_conditions = 0

        # spatial queries
        if spatial is not None:
            query = query.where(
                fn.ST_Intersects(
                    cls.geom, fn.ST_GeomFromGeoJSON(spatial.to_geometry_json())
                )
            )
            n_conditions += 1

        # temporal queries for captured time
        if temporal is not None:
            start_datetime, end_datetime = temporal
            if start_datetime is not None:
                query = query.where(cls.start_datetime >= start_datetime)
                n_conditions += 1
            if end_datetime is not None:
                query = query.where(cls.end_datetime < end_datetime)
                n_conditions += 1

        return query, n_conditions

    @property
    def epsg_no(self) -> int:
        """EPSG code of the STAC item."""
        try:
            return int(self.info["properties"]["proj:epsg"])
        except Exception:
            try:
                if "paii-misc:mgrs_tile" in self.info["properties"]:
                    epsg_no = mgrs_tile_to_epsg(
                        self.info["properties"]["paii-misc:mgrs_tile"]
                    )
                elif "paii-misc:mgrs_gzd" in self.info["properties"]:
                    epsg_no = mgrs_tile_to_epsg(
                        self.info["properties"]["paii-misc:mgrs_gzd"]
                    )
                return epsg_no
            except Exception:
                raise RuntimeError(
                    f"cannot determine `epsg_no` for {self.uid} under {self.collection_d}. Please contact admin to fix."
                )

    def choose_assets(
        self,
        lst_assets: List[str],
        dest_nodata: Union[Optional[numeric], List[Optional[numeric]]] = 0,
    ) -> VRTScene:
        """Create a virtual raster scene from the STAC item, by choosing assets from the given list of assets.

        Parameters
        ----------
        lst_assets : List[str]
            A list of assets to be used.
        dest_nodata : Union[Optional[numeric], List[Optional[numeric]]], optional
            Nodata value(s) to be used, by default 0

        Returns
        -------
        VRTScene
            The created virtual raster scene.
        """
        selected_assets = OrderedDict()
        for asset_name in lst_assets:
            selected_assets[asset_name] = interpret_protocol(
                self.info["assets"][asset_name]["href"]
            )
        # TODO, later dest_nodata could be None,
        # then here we should read from the scene model's definition to get the default source nodata value
        # and determine the dest nodata value as we want
        # also, notice VRTScene.from_assets has more control params exposed,
        # we should use those to have more precise control over vrt generation
        return VRTScene.from_assets(selected_assets, dest_nodata=dest_nodata)


def extract_base_info(info_item):
    """Extract the base info required for a STAC Item from the given information dictionary."""
    uid = info_item["id"]
    acquired_start = datetimestr_to_datetime(
        info_item["properties"]["datetime"]["start_datetime"]
    )
    acquired_end = datetimestr_to_datetime(
        info_item["properties"]["datetime"]["end_datetime"]
    )

    processed = datetimestr_to_datetime(info_item["properties"]["paii-lc:processed"])
    published = datetimestr_to_datetime(info_item["properties"]["published"])
    indexed = datetimestr_to_datetime(info_item["properties"]["paii-lc:indexed"])
    transferred = datetimestr_to_datetime(
        info_item["properties"]["paii-lc:transferred"]
    )
    archived = datetimestr_to_datetime(info_item["properties"]["paii-lc:archived"])
    created = datetimestr_to_datetime(info_item["properties"]["created"])
    updated = datetimestr_to_datetime(info_item["properties"]["updated"])

    info = info_item
    geom = json.dumps(info_item["geometry"])

    return {
        "uid": uid,
        "start_datetime": acquired_start,
        "end_datetime": acquired_end,
        "processed": processed,
        "published": published,
        "indexed": indexed,
        "transferred": transferred,
        "archived": archived,
        "created": created,
        "updated": updated,
        "valid_until": None,
        "info": info,
        "geom": geom,
    }
